from google.appengine.ext import db

class Query(db.Query):
  """ A subclass of Query with a few nifty add-ons.
  
  Retains all existing functionality of the superclass (if not exact usage)
  
  """
  def fetch(self, *args, **kwargs):
    """ If no args, fetch just first value, returns None if doesn't exist
    
    >>> class QueryTestObject(db.Model):
    ...   def __str__(self): return 'Query Test Object'
    ...   def __repr__(self): return 'Query Test Object'
    >>> db.delete(QueryTestObject.all())
    >>> print Query(QueryTestObject).fetch()
    None
    
    >>> test_object = QueryTestObject()
    >>> k = test_object.put()
    >>> print Query(QueryTestObject).fetch()
    Query Test Object
        
    If args or kwargs passed, then use superclass functionality
    >>> Query(QueryTestObject).fetch(1)
    [Query Test Object]
    
    Clean up
    >>> test_object.delete()
    
    """
    fetch_func = super(Query, self).fetch
    
    if args or kwargs:
      return fetch_func(*args, **kwargs)
    
    # No args, get first 
    try:
      return fetch_func(1)[0]
    except IndexError:
      return None
  
  def filter_by_keyword(self, **kwargs):
    """ Uses Django like keyword argument for filter -- only for equality 
    
    >>> class QueryTestObject(db.Model):
    ...   first_name = db.StringProperty()
    ...   last_name = db.StringProperty()
    ...   def __str__(self): return self.first_name + ' ' + self.last_name
    ...   def __repr__(self): return self.__str__()
    >>> db.delete(QueryTestObject.all())
    
    >>> k1 = QueryTestObject(first_name="Joe", last_name='Schmoe').put()
    >>> k2 = QueryTestObject(first_name="Joe", last_name='Doe').put()
    >>> k3 = QueryTestObject(first_name="John", last_name='Doe').put()
    
    >>> Query(QueryTestObject).filter(last_name='Doe').fetch(10)
    [Joe Doe, John Doe]
    
    >>> q = Query(QueryTestObject).filter(first_name='John', last_name='Doe')
    >>> q.fetch(10)
    [John Doe]
    
    """
    for key, value in kwargs.items():
      self.filter(key + ' =', value)
    return self
  
  def filter(self, *args, **kwargs):
    """ Tweaked version of original filter function
    
    Allows Django-like keyword arguments
    Supports custom values if property supports it
    
    >>> class CustomProperty(db.StringProperty): pass
    
    >>> class QueryTestObject(db.Model):
    ...   name = CustomProperty()
    ...   def __str__(self): return self.name
    ...   def __repr__(self): return self.name
    >>> db.delete(QueryTestObject.all())
    
    >>> k1 = QueryTestObject(name="Joe").put()
    >>> k2 = QueryTestObject(name="Bob").put()
    
    Without make_value_for_query 
    
    >>> Query(QueryTestObject).filter('name = ', 'Joe').fetch(10)
    [Joe]
    
    Now using make_value_for_query
    
    >>> QueryTestObject.name.make_value_for_query = (lambda *args : 'Bob')
    >>> Query(QueryTestObject).filter('name = ', 'Joe').fetch(10)
    [Bob]
    
    Clean up
    >>> db.delete(QueryTestObject.all())
    
    """
    # Django notation
    if kwargs: return self.filter_by_keyword(**kwargs)
    
    # Default App Engine notation
    property_operator, value = args
    
    # Check if property has any modifiers
    prop = property_operator.strip().split()[0]
    prop_obj = getattr(self._model_class, prop, None)
    if prop_obj and hasattr(prop_obj, 'make_value_for_query'):
      value = prop_obj.make_value_for_query(value)
    
    return super(Query, self).filter(property_operator, value)
  
  def __len__(self):
    """ Returns count of items in query up to 1000 """
    return len(self.fetch(1000))
